import os
import nibabel as nib
from utils import model_load, mr_lax_inference

img_file_2ch = 'data/raw/mr_2ch_seg.nii.gz'
img_file_4ch = 'data/raw/mr_4ch_seg.nii.gz'
unet = model_load()
nif_sp, nif_prd, nif_prd_bp_2ch, nif_prd_bp_4ch = mr_lax_inference(unet, img_file_2ch, img_file_4ch)
save_sp = 'data/vol_sp/mr_sp.nii.gz'
nib.save(nif_sp, save_sp)
save_prd = 'data/prd/mr_sp.nii.gz'
nib.save(nif_prd, save_prd)
save_2ch_prd_bp = 'data/prd_bp/mr_2ch_prd_bp.nii.gz'
nib.save(nif_prd_bp_2ch, save_2ch_prd_bp)
save_4ch_prd_bp = 'data/prd_bp/mr_4ch_prd_bp.nii.gz'
nib.save(nif_prd_bp_4ch, save_4ch_prd_bp)
